// SPDX-License-Identifier: GPL-2.0-only
/*
 * Example code for CSV RTMR & ATTESTATION.
 *
 * Copyright (C) 2025 Hygon Info Technologies Ltd.
 *
 * Author: Liyang Han <hanliyang@hygon.cn>
 */

#include <stdio.h>
#include <stdint.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <cpuid.h>
#include <time.h>

#include "csv-guest-module-tool.h"
#include "rtmr.h"
#include "attestation.h"

/* global program name. */
char *g_prog_name;

/******************************************************************************/
/******************************* HELPER FUNCTION ******************************/

int open_csv_guest(void)
{
	int fd;

	// check dev path
	if (access(CSV_GUEST_DEV_PATH, F_OK)) {
		PR_ERR("cannot access file %s\n", CSV_GUEST_DEV_PATH);
		return -1;
	}

	fd = open(CSV_GUEST_DEV_PATH, O_RDWR);
	if (fd < 0) {
		PR_ERR("cannot open file %s\n", CSV_GUEST_DEV_PATH);
		return -1;
	}

	return fd;
}

/**
 * Return true if valid, otherwise return false.
 */
static bool _is_ascii_hexadecimal(const char c)
{
	// '0', '1', '2', ..., '9'
	if (c >= 0x30 && c <= 0x39)
		return true;
	// 'a', 'b', 'c', 'd', 'e', 'f'
	if (c >= 0x61 && c <= 0x66)
		return true;
	return false;
}

/**
 * For example, the input @str is "abcd", the data in the output @buffer will be
 * 0xab 0xcd.
 */
int convert_ascii_str_to_bin(uint8_t *buffer, int buffer_len,
			     const char *str, int str_len)
{
	char *ptr = NULL;
	uint8_t data;
	int i;

	for (i = 0; i < str_len; i += 2) {
		ptr = (char *)str + i;
		if (!_is_ascii_hexadecimal(ptr[0])) {
			PR_ERR("invalid char '%c'\n", ptr[0]);
			return -EINVAL;
		}
		if (!_is_ascii_hexadecimal(ptr[1])) {
			PR_ERR("invalid char '%c'\n", ptr[1]);
			return -EINVAL;
		}
		if (ptr[0] >= 0x61)
			data = ptr[0] - 0x61 + 10;
		else if (ptr[0] >= 0x30)
			data = ptr[0] - 0x30;
		data <<= 4;
		if (ptr[1] >= 0x61)
			data |= ((ptr[1] - 0x61) + 10);
		else if (ptr[1] >= 0x30)
			data |= (ptr[1] - 0x30);

		assert(buffer_len > (i / 2));
		buffer[i / 2] = data;
	}

	// dump binary
	PR_INFO("After convert ascii string to binary:\n");
	for (i = 0; i < buffer_len; i++) {
		if (i % 16 == 0) {
			printf("\n");
			PR_INFO("[%02d] ", i);
		}
		printf("%02x", buffer[i]);
	}
	printf("\n\n");

	return 0;
}

static int _rdseed_u64(uint64_t *val)
{
	#define RDSEED_RETRY_LOOPS	16
	bool ok;
	unsigned int retry = RDSEED_RETRY_LOOPS;

	do {
		asm volatile("rdseed %0; setc %1"
			     : "=r" (*val), "=qm" (ok));
		if (ok)
			return true;
	} while (--retry);
	return false;
}

static int _rdrand_u64(uint64_t *val)
{
	#define RDRAND_RETRY_LOOPS	10
	bool ok;
	unsigned int retry = RDRAND_RETRY_LOOPS;

	do {
		asm volatile("rdrand %0; setc %1"
			     : "=r" (*val), "=qm" (ok));
		if (ok)
			return true;
	} while (--retry);
	return false;
}

/**
 * Generate an array of random numbers.
 *
 * Use the RDSEED feature first, followed by the RDRAND feature if unavailable,
 * and fallback to C's basic random number generation code as a last resort.
 */
int gen_random_data(void *buffer, uint32_t len)
{
	uint32_t eax, ebx, ecx, edx;
	uint32_t *u32_ptr = NULL;
	int i = 0;

	assert(len % sizeof(uint64_t) == 0);

	/* Check RDSEED feature. */
	#define RDSEED_FEAT_BIT		18
	#define RDSEED_FEAT_MASK	(1U << RDSEED_FEAT_BIT)
	eax = 7;
	__get_cpuid(eax, &eax, &ebx, &ecx, &edx);
	if (ebx & RDSEED_FEAT_MASK) {
		for (i = 0; i < len / sizeof(uint64_t); i++) {
			if (_rdseed_u64((uint64_t *)buffer + i) == false)
				goto rdrand_u64_label;
		}
		PR_INFO("generate random number success. (Using RDSEED)\n");
		return 0;
	}

rdrand_u64_label:
	/* Check RDRAND feature. */
	#define RDRAND_FEAT_BIT		30
	#define RDRAND_FEAT_MASK	(1U << RDRAND_FEAT_BIT)
	eax = 1;
	__get_cpuid(eax, &eax, &ebx, &ecx, &edx);
	if(ecx & RDRAND_FEAT_MASK) {
		for (i = 0; i < len / sizeof(uint64_t); i++) {
			if (_rdrand_u64((uint64_t *)buffer + i) == false)
				goto c_rand_label;
		}
		PR_INFO("generate random number success. (Using RDRAND)\n");
		return 0;
	}

c_rand_label:
	/* C basic random number generation code. */
	srand((unsigned int)time(NULL));
	for (i = 0; i < len / sizeof(uint32_t); i++) {
		u32_ptr = (uint32_t *)buffer + i;
		*u32_ptr = rand();
	}
	PR_INFO("generate random number success. (Using C lib)\n");
	return 0;
}

void dump_bin(uint8_t *buffer, int buffer_len)
{
	int i;

	for (i = 0; i < buffer_len; i++)
		printf("%02x", buffer[i]);

	printf("\n");
}

/******************************* HELPER FUNCTION ******************************/
/******************************************************************************/

static void print_usage(void)
{
	printf("Usage:\n\n");
	printf("%s <subcommand> [options]\n\n", g_prog_name);
	printf("Subcommands: %s, %s, %s\n",
	       APP_SUBCMD_HELP, APP_SUBCMD_RTMR, APP_SUBCMD_ATTESTATION);
}

/* Entry point */
int main(int argc, char *argv[])
{
	char *subcmd = NULL;
	int sub_argc = 0;
	char **sub_argv = NULL;

	g_prog_name = argv[0];

	if (argc < 2) {
		print_usage();
		return -1;
	}

	subcmd = argv[1];
	sub_argc = argc - 1;
	sub_argv = &argv[1];

	if (!strcmp(subcmd, APP_SUBCMD_HELP)) {
		print_usage();
		return 0;
	} else if (!strcmp(subcmd, APP_SUBCMD_RTMR)) {
		return handle_subcmd_rtmr(sub_argc, sub_argv);
	} else if (!strcmp(subcmd, APP_SUBCMD_ATTESTATION)) {
		return handle_subcmd_attestation(sub_argc, sub_argv);
	} else {
		PR_ERR("Unknown subcmd: %s\n", subcmd);
		print_usage();
		return -1;
	}
}
